import torch
import tiktoken
from torch.utils.data import Dataset, DataLoader

with open("the-verdict.txt", "r", encoding="utf-8") as f:  # 打开文本文件
    raw_text = f.read()  # 读取文本内容

class GPTDatasetV1(Dataset):  # 定义GPTDatasetV1类，继承自Dataset
    def __init__(self, txt, tokenizer, max_length, stride):  # 初始化方法
        self.input_ids = []  # 初始化输入ID列表
        self.target_ids = []  # 初始化目标ID列表

        token_ids = tokenizer.encode(txt)  # 对文本进行分词，将文本分成一个一个的token，然后每个token给一个id

        for i in range(0, len(token_ids) - max_length, stride):  # 使用滑动窗口将文本分块
            input_chunk = token_ids[i:i + max_length]  # 获取输入分块
            target_chunk = token_ids[i + 1: i + max_length + 1]  # 获取目标分块
            self.input_ids.append(torch.tensor(input_chunk))  # 将输入分块转换为张量并添加到列表中
            self.target_ids.append(torch.tensor(target_chunk))  # 将目标分块转换为张量并添加到列表中

    def __len__(self):  # 返回数据集的长度
        return len(self.input_ids)  # 返回输入ID的长度

    def __getitem__(self, idx):  # 获取指定索引的数据
        return self.input_ids[idx], self.target_ids[idx]  # 返回输入和目标的张量

def create_dataloader_v1(txt, batch_size=4, max_length=256,  # 创建数据加载器
                         stride=128, shuffle=True, drop_last=True, num_workers=0):
    tokenizer = tiktoken.get_encoding("gpt2")  # 初始化分词器
    dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)  # 创建数据集
    dataloader = DataLoader(  # 创建数据加载器
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        drop_last=drop_last,  # 如果最后一个批次小于指定批次大小，则丢弃
        num_workers=num_workers  # 使用的CPU进程数量
    )
    return dataloader  # 返回数据加载器

vocab_size = 50257  # 词汇表大小
output_dim = 256  # 输出维度
token_embedding_layer = torch.nn.Embedding(vocab_size, output_dim)  # 实例化词元嵌入层

max_length = 4  # 最大长度
dataloader = create_dataloader_v1(  # 创建数据加载器
    raw_text, batch_size=8, max_length=max_length, stride=max_length, shuffle=False
)
data_iter = iter(dataloader)  # 将数据加载器转换为Python迭代器
inputs, targets = next(data_iter)  # 获取下一个批次的数据
print("Token IDs:\n", inputs)  # 打印词元ID
print("\nInputs shape:\n", inputs.shape)  # 打印输入的形状
# 上面的代码打印以下输出：
# Token IDs:
# tensor([[   40,   367,  2885,  1464],
#         [ 1807,  3619,   402,   271],
#         [10899,  2138,   257,  7026],
#         [15632,   438,  2016,   257],
#         [  922,  5891,  1576,   438],
#         [  568,   340,   373,   645],
#         [ 1049,  5975,   284,   502],
#         [  284,  3285,   326,    11]])
#
# Inputs shape:
# torch.Size([8, 4]) ，第一个维度8代表batch_size，第二个维度4代表每一个批次取到的序列长度
token_embeddings = token_embedding_layer(inputs)  # 获取词元嵌入
print(token_embeddings.shape)  # 打印词元嵌入的形状
# 上述打印函数返回以下内容：
# torch.Size([8, 4, 256]) ，第三个维度表示每个token表示成张量的维度

# 对于GPT模型的绝对嵌入方法，我们只需要创建另一个具有与token_embedding_layer相同维度的嵌入层：
context_length = max_length  # 上下文长度等于最大长度
pos_embedding_layer = torch.nn.Embedding(context_length, output_dim)  # 创建位置嵌入层
pos_embeddings = pos_embedding_layer(torch.arange(context_length))  # 获取位置嵌入
print(pos_embeddings.shape)  # 打印位置嵌入的形状,torch.Size([4, 256])
'''
如前面的代码示例所示，pos_embeddings的输入通常是占位符向量torch.arange(context_length)，
其中包含一系列数字0、1、…，直到最大输入长度-1。context_length是表示LLM支持的输入大小的变量。
在这里，我们选择它类似于输入文本的最大长度。在实践中，输入文本可能比支持的上下文长度长，在这种情况下，我们必须截断文本。
'''
input_embeddings = token_embeddings + pos_embeddings  # 将位置嵌入与词元嵌入相加
print(input_embeddings.shape)  # 打印输入嵌入的形状
# 上述打印函数返回以下内容：
# torch.Size([8, 4, 256])


